[New] [Offload] [1/2] Disambiguate synchronous `__setitem__/offload` from `update_offload` by kylesayrs · Pull Request #768 · vllm-project/compressed-tensors

kylesayrs · 2026-07-06T23:11:48Z

Background

The ability to update offloaded values is something that can be done asynchronously across ranks.
The ability to create new offloaded values is something that must be done synchronously across ranks.

Currently on main, we do both through the same interface:

# this is a creation if the shapes don't match
# and an update if the shapes do match
module.weight = torch.nn.Parameter(new_value, ...)

# when updating, we assume that the shapes always match
def update_offload_parameter():
    ...
    # triggers update iff shapes match
    cache[name] = data

This is bad because it's very unclear to users when parameters are being updated vs created, which has serious implications for rank synchronization. It's better to make it very explicit: assignment is synchronous, while update_offload_parameter is always asynchronous.

Summary

Always offload values when assigning parameters
- Explicitly the key if it already exists: better cleanup for disk offloading
Force update_offload_parameter to call update_offload explicitly
- Make this function respect onloading_disabled
Miscellaneous
- Make tensors contiguous before writing to disk, as required by safetensors
- Use torch.no_grad decorator for update_offload_parameter
- Add offload_folder fixture for easy testing
- Add test_distributed_compression_with_disk_offload for better coverage

Testing

Update low-level unit tests to reflect new behavior of assignment
https://github.com/neuralmagic/llm-compressor-testing/actions/runs/33812157889

coderabbitai · 2026-07-06T23:15:01Z

📝 Walkthrough

Walkthrough

This PR changes offload update paths to validate stored tensors explicitly, delete existing cache entries before replacement, and handle missing disk-index entries by skipping updates. Tests were updated for local, GPU, and distributed update flows.

Changes

Offload update behavior

Layer / File(s)	Summary
OffloadCache setitem rework `src/compressed_tensors/offload/cache/base.py`	Existing keys are deleted before the new value is re-offloaded and stored.
update_offload_parameter rewrite `src/compressed_tensors/offload/__init__.py`	`update_offload_parameter` now runs under `@torch.no_grad()`, updates through `cache.update_offload(...)` for offloaded values, and only copies data in the non-cache path.
DiskCache update skip `src/compressed_tensors/offload/cache/disk.py`	`update_offload` now returns early when the offloaded tensor is absent from the index, after checking the incoming data shape/device state.
Local and GPU semantics tests `tests/test_offload/test_interface.py`, `tests/test_offload/test_module.py`, `tests/test_offload/cache/test_disk.py`	Adds coverage for offload-only updates, adjusts existing assertions, and updates overwrite behavior checks and a comment.
Distributed update paths `tests/test_offload/cache/test_dist_cpu.py`, `tests/test_offload/cache/test_dist_disk.py`, `tests/test_offload/convert/test_from_accelerate.py`	Distributed tests now read the current offloaded tensor and update through `update_offload` or `update_offload_parameter` instead of direct assignment or in-place mutation.

Estimated code review effort: 4 (Complex) | ~45 minutes

Possibly related PRs

vllm-project/compressed-tensors#666: Changes DiskCache.update_offload validation and update behavior in the same file.
vllm-project/compressed-tensors#709: Touches the same offload update paths and test coverage around them.
vllm-project/compressed-tensors#767: Also modifies update_offload_parameter and OffloadCache.__setitem__.

Suggested reviewers: HDCharles, brian-dellabetta

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly identifies the main change: separating synchronous setitem/offload behavior from update_offload. The release and series markers do not make it misleading.
Description check	✅ Passed	The description directly explains the synchronous assignment and asynchronous update changes, their synchronization implications, and the related tests.

✨ Finishing Touches

📝 Generate docstrings

Create stacked PR
Commit on current branch

🧪 Generate unit tests (beta)

Create PR with unit tests
Commit unit tests in branch kylesayrs/synchronous-setitem

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

mergify · 2026-07-07T14:11:05Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @kylesayrs.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

coderabbitai

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/test_offload/test_interface.py`:
- Around line 98-111: The assertion in test_update_offload_parameter_only is
using value inequality on tensors, which can succeed for the wrong reason if
both tensors contain the same value. Update the checks around
offloaded_linear.weight and the local offload variable to assert they are
distinct by identity or by comparing device/state, using the existing test
helpers disable_offloading and disable_onloading to verify the onloaded and
offloaded tensors are not the same object.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 219c97bf-7825-4727-9f61-43365bbb801e

📥 Commits

Reviewing files that changed from the base of the PR and between a4ee36a and f05c0c9.

📒 Files selected for processing (5)

src/compressed_tensors/offload/__init__.py
src/compressed_tensors/offload/cache/base.py
tests/test_offload/cache/test_disk.py
tests/test_offload/test_interface.py
tests/test_offload/test_module.py

🔗 Linked repositories identified

CodeRabbit considers these linked repositories for cross-repo context during reviews:

vllm-project/llm-compressor (manual)

coderabbitai · 2026-07-08T19:31:31Z

+@pytest.mark.unit
+@requires_gpu
+def test_update_offload_parameter_only(offloaded_linear: torch.nn.Linear):
+    offloaded_linear.weight = torch.nn.Parameter(torch.tensor(0.0))
+
+    with disable_offloading():
+        _ = offloaded_linear.weight
+        update_offload_parameter(offloaded_linear, "weight", torch.tensor(1))
+
+        # updating offload does not update onload
+        with disable_onloading():
+            offload = offloaded_linear.weight
+        assert offload == 1
+        assert offloaded_linear.weight != offload


🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash # 1. Check device constants in test_interface.py rg -n 'ONLOAD_DEVICE|OFFLOAD_DEVICE' tests/test_offload/test_interface.py | head -10 # 2. Check CPUCache.update_offload implementation ast-grep outline src/compressed_tensors/offload/cache/cpu.py --match update_offload --view expanded # 3. If outline doesn't show the body, read the method directly rg -n -A 15 'def update_offload' src/compressed_tensors/offload/cache/cpu.py

Repository: vllm-project/compressed-tensors

Length of output: 1085

🏁 Script executed:

#!/bin/bash set -euo pipefail echo "== tests/test_offload/test_interface.py ==" sed -n '1,180p' tests/test_offload/test_interface.py echo echo "== src/compressed_tensors/offload/cache/cpu.py ==" sed -n '1,220p' src/compressed_tensors/offload/cache/cpu.py echo echo "== src/compressed_tensors/offload/interface.py ==" sed -n '1,260p' src/compressed_tensors/offload/interface.py

Repository: vllm-project/compressed-tensors

Length of output: 7714

🏁 Script executed:

#!/bin/bash set -euo pipefail rg -n -A 80 'def update_offload_parameter|keep_onloaded_values|disable_offloading|disable_onloading' src/compressed_tensors/offload tests/test_offload

Repository: vllm-project/compressed-tensors

Length of output: 50390

Use identity or device comparison here
!= compares tensor values, so this can pass/fail for the wrong reason once both tensors contain 1. Use is not (or compare .device) to assert the onloaded and offloaded tensors are distinct.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_offload/test_interface.py` around lines 98 - 111, The assertion in test_update_offload_parameter_only is using value inequality on tensors, which can succeed for the wrong reason if both tensors contain the same value. Update the checks around offloaded_linear.weight and the local offload variable to assert they are distinct by identity or by comparing device/state, using the existing test helpers disable_offloading and disable_onloading to verify the onloaded and offloaded tensors are not the same object.

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

mergify · 2026-08-17T02:57:21Z

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews

	Protection	Waiting on
🔴	Require one maintainer review	👀 reviews

🔴 Require one maintainer review

Waiting for any of

approved-reviews-by=HDCharles
approved-reviews-by=brian-dellabetta
approved-reviews-by=dsikka
approved-reviews-by=kylesayrs

This rule is failing.

All PRs must have at least one approving review from a maintainer before merging.

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

kylesayrs force-pushed the kylesayrs/synchronous-setitem branch from d94df2d to 1509629 Compare July 6, 2026 23:12

kylesayrs changed the base branch from main to kylesayrs/revert-update July 6, 2026 23:12

This was referenced Jul 7, 2026

[Tests] [Bugfix] Skip MoE tests for models which cannot be imported vllm-project/llm-compressor#2890

Merged

Revert "[Offload] [1/2] Disambiguate synchronous `__setitem__/offload… #767

Merged

Base automatically changed from kylesayrs/revert-update to main July 7, 2026 14:08

mergify Bot added the needs-rebase label Jul 7, 2026

kylesayrs force-pushed the kylesayrs/synchronous-setitem branch from 1509629 to f05c0c9 Compare July 8, 2026 19:22

coderabbitai Bot reviewed Jul 8, 2026

View reviewed changes

mergify Bot removed the needs-rebase label Jul 8, 2026

disambiguate update parameter

0e8f1e0

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

kylesayrs force-pushed the kylesayrs/synchronous-setitem branch from f9e5240 to 0e8f1e0 Compare July 8, 2026 20:49

thing

205d730

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

kylesayrs force-pushed the kylesayrs/synchronous-setitem branch from 1d6a264 to 205d730 Compare July 8, 2026 21:19

kylesayrs mentioned this pull request Sep 3, 2026

Offloading buffer #865

Open

kylesayrs added 3 commits September 3, 2026 17:23

Merge branch 'main' into kylesayrs/synchronous-setitem

5451e78

better handling for module parallel

967f5e8

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

minor cleanup

cf06ab0

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Navigation Menu

Uh oh!

[New] [Offload] [1/2] Disambiguate synchronous `setitem/offload` from `update_offload` - #768

[New] [Offload] [1/2] Disambiguate synchronous `setitem/offload` from `update_offload`#768
kylesayrs wants to merge 5 commits into
mainfrom
kylesayrs/synchronous-setitem

kylesayrs commented Jul 6, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented Jul 6, 2026 •

edited

Loading

Walkthrough

Changes

Uh oh!

mergify Bot commented Jul 7, 2026

Uh oh!

coderabbitai Bot left a comment

Uh oh!

coderabbitai Bot Jul 8, 2026

Uh oh!

mergify Bot commented Aug 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

kylesayrs commented Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Background

Summary

Testing

Uh oh!

coderabbitai Bot commented Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Uh oh!

mergify Bot commented Jul 7, 2026

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot Jul 8, 2026

Choose a reason for hiding this comment

Uh oh!

mergify Bot commented Aug 17, 2026

Merge Protections

🔴 Require one maintainer review

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

kylesayrs commented Jul 6, 2026 •

edited

Loading

coderabbitai Bot commented Jul 6, 2026 •

edited

Loading